import numpy as np
import matplotlib.pyplot as plt

import pickle

units=['K km/s']


plotdict={'HDO 241.558 GHz': {'color': 'black', 'legend': 'HDO 241 GHz','fill_color':'gray'},
         'HDO 225.535GHz': {'color': 'blue', 'legend': 'HDO 225 GHz','fill_color':'cyan'},
         'H218O_203': {'color': 'green', 'legend': 'H$_2$$^{18}$O 203 GHz','fill_color':'cyan'},
         'C$^{17}$O J=(2-1)': {'color': 'red', 'legend': 'C$^{17}$O J=(2-1)','fill_color':'magenta'},
         'H$2$CO (J=3$_{1,2}$-2$_{1,1}$)': {'color': 'brown', 'legend': 'H$_2$CO (J=3$_{1,2}$-2$_{1,1}$)'},
         'CH3OH-241.813': {'color': 'purple', 'legend': 'CH$_3$OH 241.813 GHz'},
}
#plotlines=['HDO 241.558 GHz-H218O-conv','HDO 225.535GHz-H218O-conv','H218O_203','C$^{17}$O J=(2-1)', 'CH3OH-241.813', 'H$2$CO (J=3$_{1,2}$-2$_{1,1}$)']
#Omit H2CO for brevity
plotlines=['HDO 241.558 GHz','HDO 225.535GHz','H218O_203','C$^{17}$O J=(2-1)', 'CH3OH-241.813']
maxy=0.0
for profile_units in units:
   radial_profiles=pickle.load(open("radial_profiles_"+profile_units.replace(' ','').replace('/','')+".pickle", "rb"))
   fig, ax = plt.subplots()
   for line in plotlines:
      if np.max(radial_profiles[line]['y']) > maxy:
         maxy=np.max(radial_profiles[line]['y'])
      
      ax.fill_between(radial_profiles[line]['x'], radial_profiles[line]['y']-radial_profiles[line]['dy'], radial_profiles[line]['y']+radial_profiles[line]['dy'],  lw=1.25, color=plotdict[line]['color'],alpha=0.3)     
      ax.plot(radial_profiles[line]['x'], radial_profiles[line]['y'], lw=1.25,label=plotdict[line]['legend'], color=plotdict[line]['color'])

      #ax.step(x, radial_profiles[line]['y']/maxy, where='mid', color='k', lw=1.0,label=line)

   ax.plot([0.1,0.1],[-100,1000],linestyle='dotted',color='black')
   ax.legend()
   ax.set_xlim(0.0, 1.0)
   ax.set_ylim(-5.0, maxy+maxy*0.1)
   ax.set_xlabel('Radius (arcsec)')
   if profile_units =='K':
     ylabel='Brightness Temperature'
   elif profile_units == 'Jy':
     ylabel='Flux Density'
   elif profile_units == 'Jy km/s':
     ylabel ='Flux'
   elif profile_units == 'K km/s':
     ylabel='Intensity'
   ax.set_ylabel(ylabel+' ('+profile_units+')')
   #ax.set_title('Radial Profiles')
   new_tick_locations = np.array([0.2,0.4,0.6,0.8,1.0])
   ax2 = ax.twiny()
   def tick_function(X):
       V = X*400.0
       return ["%.0f" % z for z in V]

   ax2.set_xlim(ax.get_xlim())
   ax2.set_xticks(new_tick_locations)
   ax2.set_xticklabels(tick_function(new_tick_locations),fontsize='large')
   ax2.set_xlabel("Radius (au)",fontsize='large')


   plt.savefig('Radial_Profiles_Nice_'+profile_units.replace(' ','').replace('/','')+'.png',dpi=200.0)
   plt.savefig('Radial_Profiles_Nice_'+profile_units.replace(' ','').replace('/','')+'.pdf',dpi=200.0)

